from argparse import ArgumentParser
import os

import torch
from scipy.io import wavfile
import numpy as np
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

from speaker_trainer import Model
from pytorch_lightning.callbacks import Callback

from speaker_trainer.backend import compute_eer
from speaker_trainer.utils import gaussian_filter, mean_filter, median_filter
from tqdm import tqdm

def load_wav(path):
    sample_rate, audio = wavfile.read(path)
    audio = torch.FloatTensor(audio)
    return sample_rate, audio


def gaussian_voting_defense(trials, model, epsilon, num_voting=10, score_save_path="score.txt", device="cuda"):
	res = []
	for idx, item in enumerate(tqdm(trials)):
		label, enroll_path, test_path = item
		# load wav
		sample_rate, enroll_wav = load_wav(enroll_path)
		sample_rate, test_wav = load_wav(test_path)

		# init result list for this trials item
		res_item = []
		res_item.append(label)
		if device == "cuda":
			model = model.cuda()
			enroll_wav = enroll_wav.cuda()
			test_wav = test_wav.cuda()

		model.eval()
		with torch.no_grad():
			enroll_embedding = model.extract_speaker_embedding(enroll_wav).squeeze(0)
			test_embedding = model.extract_speaker_embedding(test_wav).squeeze(0)

			# score without guassion noise
			score = enroll_embedding.dot(test_embedding.T)
			denom = torch.norm(enroll_embedding) * torch.norm(test_embedding)
			score = score/denom
			score = score.cpu().detach().numpy().tolist()
			res_item.append(score)

			# score with guassion noise
			test_wav = test_wav.unsqueeze(0).repeat(num_voting, 1)
			gaussian_nosie = torch.randn_like(test_wav) * epsilon

			test_embeddings = model.extract_speaker_embedding(test_wav+gaussian_nosie)
			enroll_embedding = enroll_embedding.unsqueeze(0)
			scores = enroll_embedding.mm(test_embeddings.T)
			denom = torch.norm(enroll_embedding) * torch.norm(test_embeddings, dim=1)
			scores = scores/denom
			scores = scores.cpu().detach().numpy().tolist()[0]
			for score in scores:
				res_item.append(score)
		res.append(res_item)

	with open(score_save_path, "w") as f:
		for idx, item in enumerate(res):
			for val in item:
				f.write(str(val) + " ")
			f.write("\n")
	return res


def gaussian_filter_defense(trials, model, std, num_voting=10, score_save_path="score.txt", device="cuda"):
	print("start to gaussian_filter defense")
	res = []
	for idx, item in enumerate(tqdm(trials)):
		label, enroll_path, test_path = item
		# load wav
		sample_rate, enroll_wav = load_wav(enroll_path)
		sample_rate, test_wav = load_wav(test_path)

		# init result list for this trials item
		res_item = []
		res_item.append(label)
		if device == "cuda":
			model = model.cuda()
			enroll_wav = enroll_wav.cuda()
			test_wav = test_wav.cuda()

		model.eval()
		with torch.no_grad():
			enroll_embedding = model.extract_speaker_embedding(enroll_wav).squeeze(0)
			test_embedding = model.extract_speaker_embedding(test_wav).squeeze(0)

			# score without guassion noise
			score = enroll_embedding.dot(test_embedding.T)
			denom = torch.norm(enroll_embedding) * torch.norm(test_embedding)
			score = score/denom
			score = score.cpu().detach().numpy().tolist()
			res_item.append(score)

			# score with guassion filter
			test_wav = mean_filter(test_wav)

			test_embeddings = model.extract_speaker_embedding(test_wav)
			enroll_embedding = enroll_embedding.unsqueeze(0)
			scores = enroll_embedding.mm(test_embeddings.T)
			denom = torch.norm(enroll_embedding) * torch.norm(test_embeddings, dim=1)
			scores = scores/denom
			scores = scores.cpu().detach().numpy().tolist()[0]
			for score in scores:
				res_item.append(score)
		res.append(res_item)

	with open(score_save_path, "w") as f:
		for idx, item in enumerate(res):
			for val in item:
				f.write(str(val) + " ")
			f.write("\n")
	return res





if __name__ == "__main__":
	# args
		parser = ArgumentParser()
		parser.add_argument('--epsilon', help='', type=int, default=5)
		parser.add_argument('--num_voting', help='', type=int, default=10)
		parser.add_argument('--avd_trial_path', help='', type=str, default="trials.lst")
		parser.add_argument('--device', help='', type=str, default="cuda")
		parser = Trainer.add_argparse_args(parser)
		parser = Model.add_model_specific_args(parser)
		args = parser.parse_args()

		model = Model(**vars(args))

		# 0. pop loss Function parameter
		state_dict = torch.load(args.checkpoint_path, map_location="cpu")["state_dict"]
		loss_weights = []
		for key, value in state_dict.items():
			if "loss" in key:
				loss_weights.append(key)
		for item in loss_weights:
			state_dict.pop(item)

		# load speaker encoder state dict and trials
		model.load_state_dict(state_dict, strict=False)
		trials = np.loadtxt(args.avd_trial_path, str)

		# performance eval
		# model.cosine_evalute(trials)

		# voting
		gaussian_voting_defense(trials, model, args.epsilon, args.num_voting, args.score_save_path)

		# filter
		#gaussian_filter_defense(trials, model, args.epsilon, args.num_voting, args.score_save_path)


